# encoding: utf-8

from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.stattools import adfuller
from statsmodels.stats.diagnostic import acorr_ljungbox
from statsmodels.tsa.arima_model import ARIMA
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


def arma(X, colname):
    diff = 1
    # D = X.diff(diff).dropna()
    # D.columns = [colname]
    # D.plot()
    # print()
    # plot_acf(D)
    # plot_pacf(D)

    # 采用ADF方法进行平稳性检验
    adf = adfuller(X)
    print("差分序列的ADF 检验结果为：", adf)
    
    # if adf[1]>0.05:
    #     print (f'原始序列经检验不平稳，p值为: {adf[1]}')
    # while adf[1]>0.05:
    #     print(diff)
    #     adf=adfuller(X.diff(diff).dropna())
    #     diff += 1
    # print (f'原始序列经过{diff}阶差分后归于平稳，p值为: {adf[1]}')
    # input()

    # print("差分序列白噪声：", acorr_ljungbox(D, lags=1))
    # input()
    
    # 对模型进行定阶
    pmax = int(len(X)/10)
    qmax = int(len(X)/10)
    bic_matrix = []
    for p in range(pmax+1):
        temp = []
        for q in range(qmax+1):
            try:
                temp.append(ARIMA(X, (p, 1, q)).fit().bic)
            except:
                temp.append(None)
            bic_matrix.append(temp)
    bic_matrix = pd.DataFrame(bic_matrix)
    # 展开，找到最小值位置
    p, q = bic_matrix.stack().idxmin()
    print(F"BIC最小的p值和q值：{p},{q}")
    model = ARIMA(X, (p,1,q)).fit()
    # model.summary()
    res = model.forecast(3)
    # print(res)
    return res[0]
    